# VideoThumbnailer - Generates thumbnails from video files
# Copyright (C) 2007 Paul Hummer
# 
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

__license__ = "GPLv2"
__copyright__ = "2007, Paul Hummer"
__author__ = "Paul Hummer <paul@ironlionsoftware.no.spam.com>"

# I think it's important to note that the methodology of this code (and 
# some actual EXACT snippets of code) were based on Elisa (http://elisa.fluendo.com/)

import os
import sys
from threading import Event
try:
    import hashlib
except ImportError:
    import md5

try:
    import gobject
    import pygtk
    import gtk
except ImportError:
    print "The VideoThumbnailer requires the gtk python library"
    sys.exit(1)

try:
    import pygst
    pygst.require('0.10')
    import gst
except ImportError:
    print "The VideoThumbnailer requires the gtk gstreamer python library"
    sys.exit(1)

try:
    import Image
    import ImageStat
except ImportError:
    print "The VideoThumbnailer requires the python imaging library"
    sys.exit(1)

import thumbnailer

class VideoThumbnailerException(Exception): pass

class VideoThumbnailer(thumbnailer.Thumbnailer):
    ''' VideoThumbnailer does exactly what it's name demonstrates it does:
    create thumbnailed images of the videos.'''

    class VideoSinkBin(gst.Bin):
    
        def __init__(self, needed_caps):
            self.reset()
            gst.Bin.__init__(self)
            self._capsfilter = gst.element_factory_make('capsfilter', 'capsfilter')
    
            self.set_caps(needed_caps)
            self.add(self._capsfilter)
    
            fakesink = gst.element_factory_make('fakesink', 'fakesink')
            fakesink.set_property("sync", False)
            self.add(fakesink)
            self._capsfilter.link(fakesink)
    
            pad = self._capsfilter.get_pad("sink")
            ghostpad = gst.GhostPad("sink", pad)
    
            pad2probe = fakesink.get_pad("sink")
            pad2probe.add_buffer_probe(self.buffer_probe)
    
            self.add_pad(ghostpad)
            self.sink = self._capsfilter
    
        def set_current_frame(self, value):
            self._current_frame = value
    
        def set_caps(self, caps):
            gst_caps = gst.caps_from_string(caps)
            self._capsfilter.set_property("caps", gst_caps)
    
        def get_current_frame(self):
            frame = self._current_frame
            self._current_frame = None
            return frame
    
        def buffer_probe(self, pad, buffer):
            caps = buffer.caps
            if caps != None:
                s = caps[0]
                self.width = s['width']
                self.height = s['height']
            if self.width != None and self.height != None and buffer != None:
                self.set_current_frame(buffer.data)
            return True
    
        def reset(self):
            self.width = None
            self.height = None
            self.set_current_frame(None)
            

    def __init__(self, filename, src="video"):

        thumbnailer.Thumbnailer.__init__(self, filename, src)
        self._fileuri = 'file://%s' % (self.filename)
        
        #Initialize and use the gstreamer pipeline
        self._pipeline = gst.element_factory_make('playbin', 'playbin')
        self._sink = self.VideoSinkBin('video/x-raw-rgb,bpp=24,depth=24')
        self._blocker = Event()
        self._pipeline.set_property('video-sink', self._sink)
        self._pipeline.set_property('volume', 0)


        self.BORING_THRESHOLD = 2000
        self.HOLES_SIZE = (9, 35)
        self.HOLES_DATA = '\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x9a\x9a\x9a\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x91\x91\x91\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\xff\xff\xff\xa6\x91\x91\x91\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6\x00\x00\x00\xa6'
        
    def add_holes(self, img):
        holes = Image.fromstring('RGBA', self.HOLES_SIZE, self.HOLES_DATA)
        holes_h = holes.size[1]
        remain = img.size[1] % holes_h

        i = 0
        nbands = 0
        while i < (img.size[1] - remain):
            left_box = (0, i, holes.size[0], (nbands+1) * holes.size[1])
            img.paste(holes, left_box)

            right_box = (img.size[0] - holes.size[0], i,
                         img.size[0], (nbands+1) * holes.size[1])
            img.paste(holes, right_box)

            i += holes_h
            nbands += 1

        remain_holes = holes.crop((0, 0, holes.size[0], remain))
        remain_holes.load()
        img.paste(remain_holes, (0, i, holes.size[0], img.size[1]))
        img.paste(remain_holes, (img.size[0] - holes.size[0], i,
                                 img.size[0], img.size[1]))
        return img
    
    def interesting_image(self, img):
        '''
        Checks an image to see if it has the characteristics of an 
        'interesting' image, i.e. whether or not the image is worth 
        examining for thumbnailing
        '''
        stat = ImageStat.Stat(img)
        return True in [ i > self.BORING_THRESHOLD for i in stat.var ]

    def set_pipeline_state(self, pipeline, state):
        '''
        Does exactly what it claims : sets the state of the pipeline, and returns
        the boolean value of whether or not is successfully changed state
        '''
        status = pipeline.set_state(state)
        if status == gst.STATE_CHANGE_ASYNC:

            result = [False]
            max_try = 100
            nb_try = 0
            while not result[0] == gst.STATE_CHANGE_SUCCESS:
                if nb_try > max_try:
                    #State change failed
                    return False
                nb_try += 1
                result = pipeline.get_state(50*gst.MSECOND)

            return True
        elif status == gst.STATE_CHANGE_SUCCESS:
            return True
        else:
            return False


    def create_thumbnail(self):
        '''
        Creates the new thumbnail based on the information provided through the
        gstreamer API.  The file is an md5 of the video's file name, followed 
        by the .jpg extension
        '''

        if os.path.exists(self._thumb_file):
          return

        self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
        self._pipeline.set_property('uri', self._fileuri)

        if not self.set_pipeline_state(self._pipeline, gst.STATE_PAUSED):
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            raise VideoThumbnailerException('Cannot start the pipeline')

        if self._sink.width == None or self._sink.height == None:
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            raise VideoThumbnailerException('Unable to determine media size')
        sink_size = (self._sink.width, self._sink.height)

        try:
            duration, format = self._pipeline.query_duration(gst.FORMAT_TIME)
        except Exception, e:
        ## FIXME: precise this exception
            #Gstreamer cannot determine the media duration using playing-thumbnailing for file
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            img = self._play_for_thumb(sink_size, size, 0)
            #play found img
            if img:
                img.save(self._thumb_file)
                return
        else:
            duration /= gst.NSECOND
            try:
               img = self._seek_for_thumb(duration, sink_size) 
               #Seek found image
               if img:
                   img.save(self._thumb_file)
                   return
            except VideoThumbnailerException, e:
                # Fallback: No Image found in seek_for, falling back to play_for_thumb
                self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
                img = self._play_for_thumb(sink_size, duration)
                #Fallback-Play found img
                if img:
                    img.save(self._thumb_file)
                    return

        self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
        raise VideoThumbnailerException('Unable to create thumbnail.  Please file a bug')


    def _play_for_thumb(self, sink_size, duration=0):
        '''
        Plays the video file to gather information for generating a thumbnail
        '''
        id = None
        self._img = None

        if duration >= 250000:
            self._every = 25
        elif duration >= 200000:
            self._every = 15
        elif duration >= 10000:
            self._every = 10
        elif duration >= 5000:
            self._every = 5
        else:
            self._every = 1

        #Setting every-frame to self._every

        self._every_co = self._every

        ## How often Proceed?
        self._counter = 5
        
        def buffer_probe(pad, buffer):
            ## Proceed only every 5th frame!
            if self._every_co < self._every:
                self._every_co += 1
                return
            self._every_co = 0

            try: 
                img = Image.frombuffer("RGB", sink_size, buffer,
                                                "raw", "RGB",0, 1)
            except Exception, e:
                #Invalid frame
                pass
            else:
                self._img = img
                if self.interesting_image(self._img):

                    self._img.thumbnail((self.MAX_SIZE, self.MAX_SIZE), Image.BILINEAR)
                    if img.mode != 'RGBA':
                        img = img.convert(mode='RGBA')

                    #Releasing self._img
                    self._sink.reset()
                    pad.remove_buffer_probe(id)
                    self._blocker.set()
                    return

            self._counter -= 1
            if self._counter <= 0:
                #Counter off, resetting blocker
                # Is it better to return no image instead of a 'boring' one?
                if self._img:
                    self._img.thumbnail((size, size), Image.BILINEAR)
                    if img.mode != 'RGBA':
                        img = img.convert(mode='RGBA')

                self._sink.reset()
                pad.remove_buffer_probe(id)
                self._blocker.set()


        self.set_state_blocking(self._pipeline, gst.STATE_PLAYING)

        pad = self._sink.get_pad('sink')
        id = pad.add_buffer_probe(buffer_probe)
        self._blocker.wait()
        self._pipeline.set_state(gst.STATE_NULL)
        return self._img

    def _seek_for_thumb(self, duration, sink_size):
        '''
        Seeks through the video file to gather information for generating a thumbnail 
        '''
        frame_locations = [ 1.0 / 3.0, 2.0 / 3.0, 0.1, 0.9, 0.5 ]

        for location in frame_locations:
            abs_location = int(location * duration)

            if abs_location == 0:
                raise VideoThumbnailerException(self._filename, 'Video has a size of zero')
            
            event = self._pipeline.seek(1.0, gst.FORMAT_TIME,
                                        gst.SEEK_FLAG_FLUSH | gst.SEEK_FLAG_KEY_UNIT,
                                        gst.SEEK_TYPE_SET, abs_location,
                                        gst.SEEK_TYPE_NONE, 0)
            if not event:
                raise VideoThumbnailerException(self._filename, 'Unable to seek through video')

            if not self.set_pipeline_state(self._pipeline, gst.STATE_PAUSED):
                raise VideoThumbnailerException(self._filename, 'Unable to pause video')

            frame = self._sink.get_current_frame()

            try:
                img = Image.frombuffer("RGB", sink_size, frame, "raw", "RGB", 0, 1)
            except:
                #Invalid frame
                continue

            if self.interesting_image(img):
                break
            else:
                pass

        self._sink.reset()

        if img:
            img.thumbnail((self.MAX_SIZE, self.MAX_SIZE), Image.BILINEAR)
            if img.mode != 'RGBA':
                img = img.convert(mode='RGBA')
            self.set_pipeline_state(self._pipeline, gst.STATE_NULL)
            return img

gobject.type_register(VideoThumbnailer.VideoSinkBin)
